import os
import threading
from tkinter import Tk, Button, Label, filedialog, StringVar, Text, END, DISABLED, NORMAL, Entry, Frame

from PIL import Image, ImageOps, ImageEnhance, ImageFilter
import numpy as np


def compute_otsu_threshold(gray_array: np.ndarray) -> int:
    """Compute Otsu threshold for a grayscale uint8 image array."""
    # Histogram for 256 levels
    hist, _ = np.histogram(gray_array, bins=256, range=(0, 256))
    total = gray_array.size

    sum_total = np.dot(np.arange(256), hist)
    sum_b = 0.0
    w_b = 0
    max_var = 0.0
    threshold = 0

    for t in range(256):
        w_b += hist[t]
        if w_b == 0:
            continue
        w_f = total - w_b
        if w_f == 0:
            break
        sum_b += t * hist[t]
        m_b = sum_b / w_b
        m_f = (sum_total - sum_b) / w_f
        # Between-class variance
        var_between = w_b * w_f * (m_b - m_f) * (m_b - m_f)
        if var_between > max_var:
            max_var = var_between
            threshold = t
    return int(threshold)


def remove_border_lines(img_array: np.ndarray, border_pixels: int) -> np.ndarray:
    """删除图片边缘指定像素范围内的所有内容，设置为白色(255)。
    
    Args:
        img_array: 二值化图片数组 (0=黑, 255=白)
        border_pixels: 四周要清除的像素数
    """
    result = img_array.copy()
    h, w = result.shape
    
    if border_pixels > 0:
        # 清除上边缘
        result[:border_pixels, :] = 255
        # 清除下边缘
        result[h-border_pixels:, :] = 255
        # 清除左边缘
        result[:, :border_pixels] = 255
        # 清除右边缘
        result[:, w-border_pixels:] = 255
    
    return result


def process_image_to_bw(input_path: str,
                        border_pixels: int = 0) -> Image.Image:
    """Load an image and convert to high-contrast black text on white background (binary)."""
    with Image.open(input_path) as img:
        # 转换为灰度图
        gray = ImageOps.grayscale(img.convert("RGB"))

        # Denoise lightly before thresholding
        gray = gray.filter(ImageFilter.MedianFilter(size=3))

        # Stretch contrast to use full dynamic range
        gray = ImageOps.autocontrast(gray)

        # Further boost contrast and sharpness
        # 提升但不过度，减少文字毛刺
        gray = ImageEnhance.Contrast(gray).enhance(1.8)
        gray = ImageEnhance.Sharpness(gray).enhance(1.4)

        # Compute Otsu threshold
        arr = np.asarray(gray, dtype=np.uint8)
        thr = compute_otsu_threshold(arr)

        # Binarize
        bw = gray.point(lambda p: 255 if p >= thr else 0, mode="1")

        # Morphological smoothing to reduce jagged edges on strokes
        # Convert to L for morphology ops using min/max/rank filters
        bw = bw.convert("L")
        # Closing: dilate then erode (fills tiny gaps, smooths edges)
        bw = bw.filter(ImageFilter.MaxFilter(size=3))
        bw = bw.filter(ImageFilter.MinFilter(size=3))
        # Majority filter (rank=5 for 3x3) to remove burrs and unify edges
        try:
            bw = bw.filter(ImageFilter.RankFilter(size=3, rank=5))
        except Exception:
            # Fallback: ModeFilter if RankFilter unsupported
            bw = bw.filter(ImageFilter.ModeFilter(size=3))

        # Ensure final output is strictly black/white
        bw = bw.point(lambda p: 255 if p >= 128 else 0, mode="1")

        # 转换为numpy数组进行进一步处理
        arr = np.asarray(bw.convert("L"), dtype=np.uint8)
        
        # 删除边缘边框
        if border_pixels > 0:
            arr = remove_border_lines(arr, border_pixels)
        
        # 转换回PIL Image
        bw = Image.fromarray(arr, mode="L")
        bw = bw.point(lambda p: 255 if p >= 128 else 0, mode="1")

        return bw


def ensure_output_dir(selected_dir: str) -> str:
    base = os.path.basename(os.path.normpath(selected_dir))
    parent = os.path.dirname(os.path.normpath(selected_dir))
    out_dir = os.path.join(parent, f"{base}_processed")
    os.makedirs(out_dir, exist_ok=True)
    return out_dir


def is_image_file(filename: str) -> bool:
    ext = os.path.splitext(filename)[1].lower()
    return ext in {".png", ".jpg", ".jpeg", ".bmp", ".tif", ".tiff", ".webp"}


class App:
    def __init__(self) -> None:
        self.root = Tk()
        self.root.title("批量图片二值化工具（黑字白底/增强/去噪/去边框）")
        self.root.geometry("780x650")

        self.selected_dir_var = StringVar(value="未选择")

        self.lbl_dir = Label(self.root, textvariable=self.selected_dir_var, anchor="w", wraplength=750)
        self.lbl_dir.pack(padx=12, pady=(12, 6), fill="x")

        self.btn_choose = Button(self.root, text="选择文件夹", command=self.choose_dir)
        self.btn_choose.pack(padx=12, pady=6, anchor="w")

        # 参数设置区域
        settings_frame = Frame(self.root)
        settings_frame.pack(padx=12, pady=6, fill="x")
        
        # 边框清除设置
        border_frame = Frame(settings_frame)
        border_frame.pack(fill="x", pady=(0, 6))
        Label(border_frame, text="边缘清除像素数:").pack(side="left", padx=(0, 5))
        self.border_entry = Entry(border_frame, width=10)
        self.border_entry.insert(0, "0")
        self.border_entry.pack(side="left")
        Label(border_frame, text="(0=不删除边缘，建议值: 5-20)").pack(side="left", padx=(5, 0))

        self.btn_start = Button(self.root, text="开始处理", command=self.start_processing, state=DISABLED)
        self.btn_start.pack(padx=12, pady=(0, 12), anchor="w")

        self.log = Text(self.root, height=15)
        self.log.pack(padx=12, pady=6, fill="both", expand=True)

        self.selected_dir: str | None = None

    def log_line(self, message: str) -> None:
        self.log.insert(END, message + "\n")
        self.log.see(END)

    def choose_dir(self) -> None:
        d = filedialog.askdirectory()
        if d:
            self.selected_dir = d
            self.selected_dir_var.set(f"已选择：{d}")
            self.btn_start.config(state=NORMAL)
        else:
            self.selected_dir = None
            self.selected_dir_var.set("未选择")
            self.btn_start.config(state=DISABLED)

    def start_processing(self) -> None:
        if not self.selected_dir:
            self.log_line("请先选择文件夹！")
            return

        # Run in background thread to keep UI responsive
        self.btn_start.config(state=DISABLED)
        self.btn_choose.config(state=DISABLED)
        threading.Thread(target=self._process_folder_safe, daemon=True).start()

    def _process_folder_safe(self) -> None:
        try:
            self._process_folder()
        finally:
            self.btn_start.config(state=NORMAL)
            self.btn_choose.config(state=NORMAL)

    def _process_folder(self) -> None:
        assert self.selected_dir is not None
        input_dir = self.selected_dir
        output_dir = ensure_output_dir(input_dir)
        self.log_line(f"输出目录：{output_dir}")

        # 读取参数
        try:
            border_pixels = int(self.border_entry.get() or "0")
        except ValueError:
            border_pixels = 0
            self.log_line("警告：边缘清除像素数无效，使用默认值0")

        self.log_line(f"参数设置：边缘清除={border_pixels}像素")

        files = sorted(os.listdir(input_dir))
        image_files = [f for f in files if is_image_file(f)]

        if not image_files:
            self.log_line("所选文件夹内未发现图片。支持：png/jpg/jpeg/bmp/tif/tiff/webp")
            return

        total = len(image_files)
        self.log_line(f"开始处理，共 {total} 张图片……")

        processed = 0
        for idx, name in enumerate(image_files, start=1):
            src_path = os.path.join(input_dir, name)
            stem, _ = os.path.splitext(name)
            dst_path = os.path.join(output_dir, f"{stem}.png")  # 统一保存为 PNG

            try:
                bw_img = process_image_to_bw(
                    src_path,
                    border_pixels=border_pixels
                )
                bw_img.save(dst_path, format="PNG")
                processed += 1
                self.log_line(f"[{idx}/{total}] 已处理：{name} -> {os.path.basename(dst_path)}")
            except Exception as e:
                self.log_line(f"[{idx}/{total}] 跳过（读取/处理失败）：{name}，原因：{e}")

        self.log_line(f"完成。成功处理 {processed}/{total} 张图片。")

    def run(self) -> None:
        self.root.mainloop()


if __name__ == "__main__":
    App().run()


